from collections import defaultdict, Counter

# ================================核心代码段=========================================
def get_status(vocab):
    """
    核心循环中的EXAMPLE:
        word = 'T h e e'
        两两组合成： [('T', 'h'), ('h', 'e'), ('e', 'e')]
    Input:
        vocab: Dict[str, int]  # vocab统计了词语出现的词频
    Output:
        pairs: Dict[Tuple[str, str], int] # 字母对，pairs统计了单词对出现的频率
    """
    pairs = defaultdict(int)        # 当key不存在时,默认val=0
    for word, freq in vocab.items():
        symbols = word.split()      # symbols: ['l', 'o', 'w']
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return pairs

def merge_vocab(pair, v_in):
    '''
    func: 合并词表中所有的所有的best_pair
    :param (best_)pair:     ('l', 'o')
    :param v_in:            Counter({'l o w': 1, 'l o w e r': 1, 'n e w e s t': 1, 'w i d e s t': 1, 'r o o m': 1, 'r o o m s': 1})
    :return:
    '''
    v_out = {}
    bigram = ''.join(pair)  # pair:('l', 'o') -> bigram: 'lo'
    for word in v_in:
        new_word = word.replace(' '.join(pair), bigram) # new_word: 'lo w' 如果param_1不存在,不会执行替换操作
        v_out[new_word] = v_in[word]    # 获取到对应的frequency数值
    return v_out

def bpe(text, num_iters):
    vocab = Counter([' '.join(word) for word in text])  # Counter({'l o w': 1, 'l o w e r': 1, 'n e w e s t': 1, 'w i d e s t': 1, 'r o o m': 1, 'r o o m s': 1})
    for i in range(num_iters):
        pairs = get_status(vocab)
        if not pairs:       # dict({('l', 'o'): 2, ('o', 'w'): 2, ('w', 'e'): 2, ('e', 'r'): 1, ('n', 'e'): 1, ('e', 'w'): 1, ('e', 's'): 2, ('s', 't'): 2, ('w', 'i'): 1, ('i', 'd'): 1, ('d', 'e'): 1, ('r', 'o'): 2, ('o', 'o'): 2, ('o', 'm'): 2, ('m', 's'): 1})
            break
        best_pair = max(pairs, key=pairs.get)   # 在字典 pairs 中找出值最大的键
        vocab = merge_vocab(best_pair, vocab)
    return vocab

# ================================核心代码段=========================================
def post_process(vocab):
    final_vocab = defaultdict(int)
    for word, freq in vocab.items():
        wl = word.split()
        for i in wl:
            final_vocab[i] += 1
    return final_vocab

if __name__=="__main__":
    text = ['low', 'lower', 'newest', 'widest', 'room', 'rooms', 'low']
    num_iters = 4
    vocab = bpe(text, num_iters)
    print(vocab)
    final_vocab = post_process(vocab)
    print(final_vocab)